{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install llama-index nltk milvus pymilvus langchain openai python-dotenv requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package stopwords to\n",
      "[nltk_data]     /Users/yujiantang/nltk_data...\n",
      "[nltk_data]   Unzipping corpora/stopwords.zip.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import nltk\n",
    "import ssl\n",
    "\n",
    "try:\n",
    "    _create_unverified_https_context = ssl._create_unverified_context\n",
    "except AttributeError:\n",
    "    pass\n",
    "else:\n",
    "    ssl._create_default_https_context = _create_unverified_https_context\n",
    "\n",
    "nltk.download(\"stopwords\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    GPTVectorStoreIndex, \n",
    "    GPTSimpleKeywordTableIndex, \n",
    "    SimpleDirectoryReader,\n",
    "    LLMPredictor,\n",
    "    ServiceContext,\n",
    "    StorageContext\n",
    ")\n",
    "from langchain.llms.openai import OpenAIChat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "    __  _________ _   ____  ______\n",
      "   /  |/  /  _/ /| | / / / / / __/\n",
      "  / /|_/ // // /_| |/ / /_/ /\\ \\\n",
      " /_/  /_/___/____/___/\\____/___/ {Lite}\n",
      "\n",
      " Welcome to use Milvus!\n",
      "\n",
      " Version:   v2.2.8-lite\n",
      " Process:   49630\n",
      " Started:   2023-05-18 16:00:02\n",
      " Config:    /Users/yujiantang/.milvus.io/milvus-server/2.2.8/configs/milvus.yaml\n",
      " Logs:      /Users/yujiantang/.milvus.io/milvus-server/2.2.8/logs\n",
      "\n",
      " Ctrl+C to exit ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[93m[has_collection] retry:4, cost: 0.27s, reason: <_MultiThreadedRendezvous: StatusCode.UNAVAILABLE, internal: Milvus Proxy is not ready yet. please wait>\u001b[0m\n",
      "\u001b[93m[has_collection] retry:5, cost: 0.81s, reason: <_MultiThreadedRendezvous: StatusCode.UNAVAILABLE, internal: Milvus Proxy is not ready yet. please wait>\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import openai\n",
    "load_dotenv()\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "from llama_index.vector_stores import MilvusVectorStore\n",
    "from milvus import default_server\n",
    "\n",
    "default_server.start()\n",
    "vector_store = MilvusVectorStore(\n",
    "   host = \"127.0.0.1\",\n",
    "   port = default_server.listen_port\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_titles = [\"Toronto\", \"Seattle\", \"San Francisco\", \"Chicago\", \"Boston\", \"Washington, D.C.\", \"Cambridge, Massachusetts\", \"Houston\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "for title in wiki_titles:\n",
    "    response = requests.get(\n",
    "        'https://en.wikipedia.org/w/api.php',\n",
    "        params={\n",
    "            'action': 'query',\n",
    "            'format': 'json',\n",
    "            'titles': title,\n",
    "            'prop': 'extracts',\n",
    "            # 'exintro': True,\n",
    "            'explaintext': True,\n",
    "        }\n",
    "    ).json()\n",
    "    page = next(iter(response['query']['pages'].values()))\n",
    "    wiki_text = page['extract']\n",
    "\n",
    "    data_path = Path('data')\n",
    "    if not data_path.exists():\n",
    "        Path.mkdir(data_path)\n",
    "\n",
    "    with open(data_path / f\"{title}.txt\", 'w') as fp:\n",
    "        fp.write(wiki_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all wiki documents\n",
    "city_docs = {}\n",
    "for wiki_title in wiki_titles:\n",
    "    city_docs[wiki_title] = SimpleDirectoryReader(input_files=[f\"data/{wiki_title}.txt\"]).load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/yujiantang/Documents/workspace/hello_world_project/hw_milvus/lib/python3.10/site-packages/langchain/llms/openai.py:696: UserWarning: You are trying to use a chat model. This way of initializing it is no longer supported. Instead, please use: `from langchain.chat_models import ChatOpenAI`\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "llm_predictor_chatgpt = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=\"gpt-3.5-turbo\"))\n",
    "service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor_chatgpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build city document index\n",
    "city_indices = {}\n",
    "index_summaries = {}\n",
    "for wiki_title in wiki_titles:\n",
    "    city_indices[wiki_title] = GPTVectorStoreIndex.from_documents(city_docs[wiki_title], service_context=service_context, storage_context=storage_context)\n",
    "    # set summary text for city\n",
    "    index_summaries[wiki_title] = f\"Wikipedia articles about {wiki_title}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.composability import ComposableGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = ComposableGraph.from_indices(\n",
    "    GPTSimpleKeywordTableIndex,\n",
    "    [index for _, index in city_indices.items()], \n",
    "    [summary for _, summary in index_summaries.items()],\n",
    "    max_keywords_per_chunk=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.query.query_transform.base import DecomposeQueryTransform\n",
    "decompose_transform = DecomposeQueryTransform(\n",
    "    llm_predictor_chatgpt, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_engine.transform_query_engine import TransformQueryEngine\n",
    "custom_query_engines = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in city_indices.values():\n",
    "    query_engine = index.as_query_engine(service_context=service_context)\n",
    "    transform_extra_info = {'index_summary': index.index_struct.summary}\n",
    "    tranformed_query_engine = TransformQueryEngine(query_engine, decompose_transform, transform_extra_info=transform_extra_info)\n",
    "    custom_query_engines[index.index_id] = tranformed_query_engine\n",
    "\n",
    "custom_query_engines[graph.root_index.index_id] = graph.root_index.as_query_engine(\n",
    "    retriever_mode='simple', \n",
    "    response_mode='tree_summarize', \n",
    "    service_context=service_context\n",
    ")\n",
    "\n",
    "query_engine_decompose = graph.as_query_engine(\n",
    "    custom_query_engines=custom_query_engines,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What is the name of the airport in Seattle?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What is the name of the airport in Seattle?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are the major airports in Houston?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are the major airports in Houston?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are some notable features of the Toronto airport?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the airports in Seattle, Houston, and Toronto. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are some notable features of the Toronto Pearson International Airport?\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response_chatgpt = query_engine_decompose.query(\n",
    "    \"Compare and contrast the airports in Seattle, Houston, and Toronto. \"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Seattle has one major airport called Seattle-Tacoma International Airport, while Houston has two major airports called George Bush Intercontinental Airport and William P. Hobby Airport, as well as a third municipal airport called Ellington Airport. Toronto's busiest airport is called Toronto Pearson International Airport and is located on the city's western boundary with Mississauga. It offers limited commercial and passenger service to nearby destinations in Canada and the United States. Seattle-Tacoma International Airport and George Bush Intercontinental Airport are both major international airports, while William P. Hobby Airport and Ellington Airport are smaller and serve more regional destinations. Toronto Pearson International Airport is Canada's busiest airport and offers a direct link to Union Station through the Union Pearson Express train service.\n"
     ]
    }
   ],
   "source": [
    "print(str(response_chatgpt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_query_engines = {}\n",
    "for index in city_indices.values():\n",
    "    query_engine = index.as_query_engine(service_context=service_context)\n",
    "    custom_query_engines[index.index_id] = query_engine\n",
    "\n",
    "custom_query_engines[graph.root_index.index_id] = graph.root_index.as_query_engine(\n",
    "    retriever_mode='simple', \n",
    "    response_mode='tree_summarize', \n",
    "    service_context=service_context\n",
    ")\n",
    "\n",
    "query_engine = graph.as_query_engine(\n",
    "    custom_query_engines=custom_query_engines,    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The context information provided does not contain enough information to answer the question.'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_chatgpt = query_engine.query(\n",
    "    \"Compare and contrast the airports in Seattle, Houston, and Toronto. \"\n",
    ")\n",
    "str(response_chatgpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the sports environment of Houston and Boston. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What sports teams are based in Houston?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the sports environment of Houston and Boston. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What sports teams are based in Houston?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the sports environment of Houston and Boston. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are some notable sports teams based in Boston?\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m> Current query: Compare and contrast the sports environment of Houston and Boston. \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3m> New query: What are some notable sports teams based in Boston?\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response_chatgpt = query_engine_decompose.query(\n",
    "    \"Compare and contrast the sports environment of Houston and Boston. \"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Houston has sports teams for every major professional league except the National Hockey League, while Boston has teams for the NFL, MLB, NBA, and NHL. Both cities have professional soccer teams, with Houston having a Major League Soccer franchise and Boston having a National Women's Soccer League team. Boston also has several college sports teams and Esports teams, while Houston does not have any notable college sports teams or Esports teams. Both cities are known for hosting major sporting events, with Boston hosting the Boston Marathon and the Head of the Charles Regatta, while Houston does not have any comparable events. Overall, Boston has a more diverse sports environment with teams in all major professional leagues and a strong presence in college sports and Esports.\""
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(response_chatgpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hw_milvus",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
